#!/usr/bin/env python3
"""
安全加固验证脚本
测试新的安全执行器是否能够有效阻止恶意代码
"""

import asyncio
import aiohttp
import time
import sys
import os
from typing import Dict, List, Any

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

class SecurityHardeningVerifier:
    """安全加固验证器"""
    
    def __init__(self, base_url: str = "http://localhost:8000"):
        self.base_url = base_url
        self.session = None
        self.test_results = []
    
    async def setup(self):
        """设置测试环境"""
        if not self.session:
            self.session = aiohttp.ClientSession()
    
    async def cleanup(self):
        """清理测试环境"""
        if self.session:
            await self.session.close()
    
    async def submit_task(self, code: str, code_type: str = "python", 
                         security_level: str = "low") -> str:
        """提交测试任务"""
        if not self.session:
            await self.setup()
            
        request_data = {
            "code": code,
            "code_type": code_type,
            "user_info": {
                "user_id": f"security_test_{int(time.time())}",
                "username": "安全测试",
                "permissions": ["execute"]
            },
            "security_level": security_level,
            "timeout_seconds": 10,
            "memory_limit_mb": 64,
            "cpu_limit": 0.5,
            "environment_vars": {}
        }
        
        assert self.session is not None
        async with self.session.post(f"{self.base_url}/execute", json=request_data) as response:
            if response.status != 200:
                raise Exception(f"任务提交失败: HTTP {response.status}")
            
            result = await response.json()
            if not result.get("success"):
                raise Exception(f"任务提交失败: {result.get('message')}")
            
            return result["data"]["task_id"]
    
    async def wait_for_completion(self, task_id: str, timeout: int = 15) -> Dict[str, Any]:
        """等待任务完成"""
        if not self.session:
            await self.setup()
            
        start_time = time.time()
        
        while time.time() - start_time < timeout:
            assert self.session is not None
            async with self.session.get(f"{self.base_url}/status/{task_id}") as response:
                if response.status == 200:
                    result = await response.json()
                    
                    if result.get("success") and result.get("data"):
                        status_data = result["data"]
                        if status_data.get("status") in ["completed", "failed", "timeout"]:
                            return status_data
            
            await asyncio.sleep(0.5)
        
        raise Exception(f"任务 {task_id} 等待超时")
    
    async def run_security_test(self, test_name: str, malicious_code: str, 
                               code_type: str = "python", 
                               expected_blocked: bool = True) -> Dict[str, Any]:
        """运行单个安全测试"""
        print(f"  🧪 测试: {test_name}")
        
        start_time = time.time()
        result = {
            "name": test_name,
            "code": malicious_code,
            "code_type": code_type,
            "expected_blocked": expected_blocked,
            "success": False,
            "blocked": False,
            "error": None,
            "output": None,
            "execution_time": 0.0
        }
        
        try:
            # 提交恶意代码
            task_id = await self.submit_task(malicious_code, code_type)
            
            # 等待执行完成
            execution_result = await self.wait_for_completion(task_id)
            
            result["output"] = execution_result.get("output", "")
            result["error"] = execution_result.get("error", "")
            result["return_code"] = execution_result.get("return_code", 0)
            result["status"] = execution_result.get("status", "unknown")
            
            # 判断是否被成功阻止
            output_msg = execution_result.get("output", "") or ""
            error_msg = execution_result.get("error", "") or ""
            combined_msg = (output_msg + " " + error_msg).lower()
            
            # 扩展的安全关键词列表，包含中文和英文
            security_keywords = [
                # 中文关键词
                "被禁用", "禁用", "不允许", "限制", "阻止", "安全", "违规", "检测到", 
                "安全检查失败", "代码执行被阻止", "保护系统安全", "沙箱", "危险",
                # bash特定的安全输出
                "🚫", "❌", "安全沙箱", "检测到以下违规操作",
                # Python错误类型（资源限制导致的安全阻止）
                "memoryerror", "memory error", "内存错误", "内存不足",
                "timeouterror", "timeout", "超时", "时间限制",
                "resourceerror", "resource", "资源限制",
                # 英文关键词  
                "importerror", "runtimeerror", "permission denied",
                "operation not permitted", "blocked", "restricted",
                "disabled", "not allowed", "security", "violation",
                "malicious", "dangerous", "forbidden", "unauthorized"
            ]
            
            if execution_result.get("status") == "failed":
                if any(keyword in combined_msg for keyword in security_keywords):
                    result["blocked"] = True
                    result["success"] = expected_blocked
                else:
                    result["blocked"] = False
                    result["success"] = not expected_blocked
            elif execution_result.get("status") == "timeout":
                # 超时也可能意味着被阻止（资源限制生效）
                result["blocked"] = True
                result["success"] = expected_blocked
            else:
                # 执行成功，但可能安全信息在输出中
                if any(keyword in combined_msg for keyword in security_keywords):
                    result["blocked"] = True
                    result["success"] = expected_blocked
                else:
                    # 执行成功且没有安全警告，检查是否应该被阻止
                    result["blocked"] = False
                    result["success"] = not expected_blocked
            
        except Exception as e:
            result["error"] = str(e)
            # 区分网络异常和安全阻止
            error_msg = str(e).lower()
            if any(net_err in error_msg for net_err in ["connect", "disconnected", "timeout", "connection", "server"]):
                # 网络异常，不算被阻止，但标记为网络错误
                result["blocked"] = False
                result["success"] = False  # 网络异常总是算测试失败
                result["network_error"] = True  # 标记为网络错误
            else:
                # 其他异常，可能是安全阻止
                result["blocked"] = True
                result["success"] = expected_blocked
        
        result["execution_time"] = time.time() - start_time
        
        if result["success"]:
            print(f"    ✅ 测试通过 ({result['execution_time']:.2f}s)")
            if result["expected_blocked"] and result["blocked"]:
                print(f"    🛡️  恶意代码被成功阻止")
            elif not result["expected_blocked"] and not result["blocked"]:
                print(f"    🟢 正常代码正常执行")
        else:
            print(f"    ❌ 测试失败 ({result['execution_time']:.2f}s)")
            if result.get("network_error"):
                print(f"    🌐 网络连接错误（服务器可能未启动）")
                print(f"    ❗ 错误: {result['error'][:200]}...")
            elif result["expected_blocked"] and not result["blocked"]:
                print(f"    💀 危险：恶意代码未被阻止！")
                print(f"    📋 状态: {result.get('status', 'unknown')}")
                if result.get('output'):
                    print(f"    📤 输出: {result['output'][:200]}...")
                if result.get('error'):
                    print(f"    ❗ 错误: {result['error'][:200]}...")
            elif not result["expected_blocked"] and result["blocked"]:
                print(f"    ⚠️  警告：正常代码被误杀")
                print(f"    📋 状态: {result.get('status', 'unknown')}")
                if result.get('output'):
                    print(f"    📤 输出: {result['output'][:200]}...")
                if result.get('error'):
                    print(f"    ❗ 错误: {result['error'][:200]}...")
        
        self.test_results.append(result)
        return result
    
    async def run_all_security_tests(self):
        """运行所有安全测试"""
        print("🔒 开始安全加固验证测试...\n")
        
        # Python恶意代码测试
        print("📝 Python 安全测试:")
        
        # 1. 文件系统访问测试
        await self.run_security_test(
            "读取敏感文件",
            "with open('/etc/passwd', 'r') as f: print(f.read())",
            "python",
            True
        )
        
        await self.run_security_test(
            "写入恶意文件",
            "with open('/tmp/malicious.txt', 'w') as f: f.write('pwned')",
            "python", 
            True
        )
        
        # 2. 网络访问测试
        await self.run_security_test(
            "网络请求",
            "import urllib.request\nresponse = urllib.request.urlopen('http://httpbin.org/ip')\nprint(response.read())",
            "python",
            True
        )
        
        await self.run_security_test(
            "Socket连接",
            "import socket\ns = socket.socket()\ns.connect(('127.0.0.1', 80))",
            "python",
            True
        )
        
        # 3. 系统命令执行
        await self.run_security_test(
            "执行系统命令",
            "import os\nos.system('whoami')",
            "python",
            True
        )
        
        await self.run_security_test(
            "subprocess执行",
            "import subprocess\nsubprocess.run(['ls', '/etc'])",
            "python",
            True
        )
        
        # 4. 危险模块导入
        await self.run_security_test(
            "导入os模块",
            "import os\nprint(os.getcwd())",
            "python",
            True
        )
        
        await self.run_security_test(
            "导入subprocess模块", 
            "import subprocess\nprint('subprocess imported')",
            "python",
            True
        )
        
        # 5. 内存/CPU攻击
        await self.run_security_test(
            "内存炸弹",
            "data = 'A' * (100 * 1024 * 1024)  # 100MB\nprint('Memory bomb')",
            "python",
            True
        )
        
        await self.run_security_test(
            "无限循环",
            "while True: pass",
            "python",
            True
        )
        
        # 6. 反射攻击
        await self.run_security_test(
            "eval执行",
            "eval('print(\"eval works\")')",
            "python",
            True
        )
        
        await self.run_security_test(
            "exec执行",
            "exec('import os; print(os.getcwd())')",
            "python",
            True
        )
        
        # JavaScript恶意代码测试
        print("\n📜 JavaScript 安全测试:")
        
        await self.run_security_test(
            "require模块导入",
            "const fs = require('fs'); console.log(fs.readFileSync('/etc/passwd', 'utf8'));",
            "javascript",
            True
        )
        
        await self.run_security_test(
            "process对象访问",
            "console.log(process.env); console.log(process.cwd());",
            "javascript",
            True
        )
        
        await self.run_security_test(
            "eval执行",
            "eval('console.log(\"eval works in js\")');",
            "javascript", 
            True
        )
        
        # Bash恶意代码测试
        print("\n🐚 Bash 安全测试:")
        
        await self.run_security_test(
            "文件删除命令",
            "rm -rf /tmp/*",
            "bash",
            True
        )
        
        await self.run_security_test(
            "网络请求命令",
            "curl http://httpbin.org/ip",
            "bash",
            True
        )
        
        await self.run_security_test(
            "系统信息收集",
            "cat /etc/passwd; ps aux; netstat -an",
            "bash",
            True
        )
        
        await self.run_security_test(
            "文件重定向攻击",
            "echo 'malicious' > /etc/passwd",
            "bash",
            True
        )
        
        await self.run_security_test(
            "命令注入攻击",
            "echo hello; rm -f /tmp/test; cat /etc/hosts",
            "bash",
            True
        )
        
        await self.run_security_test(
            "后台进程攻击",
            "sleep 100 &",
            "bash", 
            True
        )
        
        await self.run_security_test(
            "系统服务操作",
            "systemctl status sshd",
            "bash",
            True
        )
        
        await self.run_security_test(
            "危险模块操作",
            "modprobe test_module",
            "bash",
            True
        )
        
        # 正常代码测试（确保不被误杀）
        print("\n✅ 正常代码测试:")
        
        await self.run_security_test(
            "正常Python计算",
            "print('Hello, World!')\nresult = 1 + 1\nprint(f'1 + 1 = {result}')",
            "python",
            False
        )
        
        await self.run_security_test(
            "正常JavaScript输出",
            "console.log('Hello from JavaScript'); console.log(Math.PI);",
            "javascript",
            False
        )
        
        await self.run_security_test(
            "正常Bash命令",
            "echo 'Hello, Bash!'; echo 'Current timestamp:'; date",
            "bash",
            False
        )
        
        await self.run_security_test(
            "正常if条件语句",
            "if [ 5 -gt 3 ]; then echo 'Math works'; fi",
            "bash",
            False
        )
        
        await self.run_security_test(
            "正常变量操作",
            "name='World'; echo \"Hello, $name!\"",
            "bash",
            False
        )
        
        await self.run_security_test(
            "正常循环结构",
            "for i in 1 2 3; do echo \"Number: $i\"; done",
            "bash",
            False
        )
        
        await self.run_security_test(
            "正常数学运算",
            "result=$((5 + 3)); echo \"5 + 3 = $result\"",
            "bash",
            False
        )
    
    def generate_security_report(self):
        """生成安全测试报告"""
        total_tests = len(self.test_results)
        passed_tests = sum(1 for result in self.test_results if result["success"])
        
        # 分类统计（排除网络错误）
        malicious_tests = [r for r in self.test_results if r["expected_blocked"] and not r.get("network_error")]
        normal_tests = [r for r in self.test_results if not r["expected_blocked"] and not r.get("network_error")]
        
        malicious_blocked = sum(1 for r in malicious_tests if r["blocked"])
        normal_allowed = sum(1 for r in normal_tests if not r["blocked"])
        
        # 网络错误统计
        network_errors = sum(1 for r in self.test_results if r.get("network_error"))
        
        # 计算安全评分
        if len(malicious_tests) > 0:
            block_rate = malicious_blocked / len(malicious_tests) * 100
        else:
            block_rate = 100
        
        if len(normal_tests) > 0:
            false_positive_rate = (len(normal_tests) - normal_allowed) / len(normal_tests) * 100
        else:
            false_positive_rate = 0
        
        # 安全等级评估
        if block_rate >= 95 and false_positive_rate <= 10:
            security_level = "🟢 优秀"
        elif block_rate >= 80 and false_positive_rate <= 20:
            security_level = "🟡 良好" 
        elif block_rate >= 60 and false_positive_rate <= 30:
            security_level = "🟠 中等"
        else:
            security_level = "🔴 危险"
        
        print(f"\n" + "="*60)
        print("🔒 安全加固验证报告")
        print("="*60)
        print(f"总测试数量: {total_tests}")
        print(f"通过测试数量: {passed_tests}")
        print(f"测试通过率: {passed_tests/total_tests*100:.1f}%")
        print()
        print(f"恶意代码测试: {len(malicious_tests)} 个")
        print(f"成功阻止: {malicious_blocked} 个")
        print(f"阻止率: {block_rate:.1f}%")
        print()
        print(f"正常代码测试: {len(normal_tests)} 个")
        print(f"正常执行: {normal_allowed} 个")
        print(f"误杀率: {false_positive_rate:.1f}%")
        if network_errors > 0:
            print()
            print(f"网络错误: {network_errors} 个")
        print()
        print(f"安全等级: {security_level}")
        print("="*60)
        
        # 显示失败的测试
        failed_tests = [r for r in self.test_results if not r["success"]]
        if failed_tests:
            print("\n❌ 失败的测试用例:")
            for test in failed_tests:
                if test.get("network_error"):
                    status = "网络连接错误"
                elif test["expected_blocked"] and not test["blocked"]:
                    status = "未阻止"
                else:
                    status = "被误杀"
                print(f"  - {test['name']} ({test['code_type']}): {status}")
            
            # 统计网络错误
            network_errors = [r for r in failed_tests if r.get("network_error")]
            if network_errors:
                print(f"\n🌐 网络连接问题: {len(network_errors)} 个测试因服务器连接问题失败")
                print("   建议：请确保服务器正在运行并重新测试")
        
        # 显示被成功阻止的恶意代码
        blocked_malicious = [r for r in malicious_tests if r["blocked"]]
        if blocked_malicious:
            print(f"\n✅ 成功阻止的恶意代码 ({len(blocked_malicious)}/{len(malicious_tests)}):")
            for test in blocked_malicious:
                print(f"  - {test['name']} ({test['code_type']})")

async def main():
    """主函数"""
    verifier = SecurityHardeningVerifier()
    
    try:
        await verifier.setup()
        await verifier.run_all_security_tests()
        verifier.generate_security_report()
        
    except Exception as e:
        print(f"❌ 测试执行失败: {e}")
        return 1
    
    finally:
        await verifier.cleanup()
    
    return 0

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="安全加固验证测试")
    parser.add_argument("--url", default="http://localhost:8000", 
                       help="服务器URL (默认: http://localhost:8000)")
    
    args = parser.parse_args()
    
    # 更新全局URL
    SecurityHardeningVerifier.__init__.__defaults__ = (args.url,)
    
    # 运行测试
    exit_code = asyncio.run(main())
    sys.exit(exit_code) 